library(dplyr)
library(tidyr)
library(ggplot2)
library(viridis)

source("HNC_metadata_tidy.R")

sbsCOSMIC <- read.csv("Input_signature_attributions/Input_SBS_COSMIC_attributions.csv", stringsAsFactors = T, row.names = 1) %>%
  tibble::rownames_to_column("donor_id") %>%
  mutate(APOBEC = SBS2 + SBS13, .keep = "unused")

# Mutation burden
mut_burden <- data.frame(donor_id = sbsCOSMIC$donor_id, SBS_burden = rowSums(sbsCOSMIC[, -1]))

# relative burden
sigs_rel <- sbsCOSMIC %>% mutate_if(is.numeric, ~ (. / mut_burden[, 2]))

data <- data %>%
  mutate(
    tob_hpv = case_when(
      tobacco != "Never" & hpv_pos_opc == "Positive" ~ "tobacco + HPV",
      tobacco != "Never" & hpv_pos_opc == "Negative" ~ "tobacco",
      tobacco == "Never" & hpv_pos_opc == "Positive" ~ "HPV",
      tobacco == "Never" & hpv_pos_opc == "Negative" ~ "none"
    ),
    tob_hpv = factor(tob_hpv, levels = c("HPV", "tobacco + HPV", "none", "tobacco")),
    hpv_pos_opc = factor(hpv_pos_opc, levels = c("Positive", "Negative"))
  ) %>%
  left_join(sigs_rel, by = "donor_id")

plot_kw_rel <- function(dat, signatures, factors,
                        pval_cutoff = 0.05, output.kw = data.frame(),
                        y_labs = "Mutation burden", x_labs = f) {
  plots <- list()

  # Loop over each signature and factor
  for (s in signatures) {
    for (f in factors) {
      # Check if the factor has less than 2 categories, if so, print a warning and skip
      if (dat[, f] %>% unique() %>% length() < 2) {
        print(paste("Warning: Variable", f, "has been removed due to <2 categories"))
      } else {
        kw <- tryCatch(
          {
            kruskal.test(dat[, s], dat[, f])
          },
          error = function(e) {
            NULL
          }
        )

        # If there was an error in computing the test, move to the next factor
        if (is.null(kw)) {
          print(paste("Variable", f, "has been removed due to ERROR in KW test"))
          next
        }

        result.var <- data.frame(
          signature = paste(s),
          factor = paste(f),
          pval = kw$p.value
        )

        output.kw <- rbind(output.kw, result.var)

        # If p-value is less than the cutoff, proceed with plotting
        if (kw$p.value < pval_cutoff) {
          # Remove NAs
          subset <- dat[!is.na(dat[f]) & !is.na(dat[s]), ]

          # plot
          p_top <- ggplot(data = subset, aes(x = get(f), y = get(s), fill = get(f))) +
            geom_violin(aes(color = get(f)), width = 1, alpha = .8, linewidth = .3) +
            geom_boxplot(color = "#585858", alpha = 0.15, linewidth = .3, staplewidth = .2, outlier.shape = NA) +
            geom_jitter(fill = "#585858", size = 0.8, alpha = .25, stroke = 0, width = .1, height = .1) +
            labs(
              title = s,
              subtitle = paste("Kruskal-Wallis, p =", signif(kw$p.value, 3)),
              # subtitle = ifelse(nrow(propdata) == 2,paste("Wilcoxon, p =",signif(kw$p.value,3)),paste("Kruskal-Wallis, p =",signif(kw$p.value,3))),
              y = y_labs,
              x = str_to_title(str_replace_all(x_labs, "_", " "))
            ) +
            coord_cartesian(ylim = c(0, 1)) +
            scale_fill_viridis(discrete = TRUE, end = .8) +
            scale_color_viridis(discrete = TRUE, end = .8) +
            theme_minimal() +
            theme(
              legend.position = "none",
              plot.title = element_text(face = "bold", size = 12, hjust = 0.5),
              plot.subtitle = element_text(size = 10, hjust = 0.5),
              axis.text = element_text(size = 12),
              axis.text.x = element_text(angle = 45, hjust = 1),
              axis.line = element_line(size = .3,color = "#404040"),
              axis.ticks = element_line(size = .3,color = "#404040"),
              panel.grid = element_blank()
              # axis.title.x = element_blank(),
            )
          print(p_top)
          plots[[paste(s, f, sep = "_")]] <- p_top
        }
      }
    }
  }

  return(list("output.kw" = output.kw, "plots" = plots))
}

for (f in c("hpv_pos_opc", "tob_hpv")) {
  plot_kw_rel(data,
    signatures = "APOBEC", factors = f,
    pval_cutoff = 1, y_labs = "Relative burden"
  )
  ggsave(paste0("output/ExtendedDataFigure_7ab_", f, "_", Sys.Date(), ".pdf"),
    device = "pdf",
    width = 3.5, height = 4, dpi = 700
  )
}
